#include "Depth/DrivableSeg.h"

DrivableSeg::DrivableSeg(ros::NodeHandle &nh) {

    nh.param("N_SCAN", N_SCAN,120);
    nh.param("Horizon_SCAN", Horizon_SCAN,240);
    nh.param("ang_res_x", ang_res_x,140.0/Horizon_SCAN);
    nh.param("ang_res_y", ang_res_y,70.0/N_SCAN);
    nh.param("ang_bottom", ang_bottom,30.0);
    nh.param("vertical_bottom", vertical_bottom,20.0-0.1);
    nh.param("horizon_bottom", horizon_bottom,20.0);
    nh.param("CAR_HEIGHT", CAR_HEIGHT,0.408);

    nh.param("sensorMinimumRange", sensorMinimumRange,0.1);
    nh.param("segmentAlphaX", segmentAlphaX,ang_res_x/180*M_PI);
    nh.param("segmentAlphaY", segmentAlphaY,ang_res_y/180*M_PI);
    nh.param("segmentTheta", segmentTheta,10 / 180.0 * M_PI);
    nh.param("segmentValidPointNum", segmentValidPointNum,5);
    nh.param("segmentValidLineNum", segmentValidLineNum,5);

    nh.param("toler",toler, 0.65);
    nh.param("min_size", min_size,1);
    nh.param("max_size", max_size,100);

    row_angle_sin.clear();
    row_angle_cos.clear();

    float step_row = (float) ang_res_y / 180 * M_PI;
    float start_row = ((float) ang_bottom / 180 * M_PI);
    for (int i = 0; i < N_SCAN; i++) {
        row_angle_sin.push_back(sin(start_row));
        row_angle_cos.push_back(cos(start_row));
        start_row -= step_row;
    }


    x_mat = cv::Mat::zeros(N_SCAN, Horizon_SCAN, CV_32F);
    y_mat = cv::Mat::zeros(N_SCAN, Horizon_SCAN, CV_32F);

    nanPoint.x = std::numeric_limits<float>::quiet_NaN();
    nanPoint.y = std::numeric_limits<float>::quiet_NaN();
    nanPoint.z = std::numeric_limits<float>::quiet_NaN();
    nanPoint.intensity = -1;

    fullPointCloud.reset(new pcl::PointCloud<pcl::PointXYZI>());
    SegObjectCloud.reset(new pcl::PointCloud<pcl::PointXYZI>());
    GroundCloud.reset(new pcl::PointCloud<pcl::PointXYZI>());
    ClusteredCloudCenter.reset(new pcl::PointCloud<pcl::PointXYZ>());

    fullPointCloud->width = N_SCAN * Horizon_SCAN;
    fullPointCloud->height = 1;
    fullPointCloud->points.resize(N_SCAN * Horizon_SCAN);

    std::pair <int8_t, int8_t> neighbor;
    neighbor.first = -1;
    neighbor.second = 0;
    neighborIterator.push_back(neighbor);
    neighbor.first = 0;
    neighbor.second = 1;
    neighborIterator.push_back(neighbor);
    neighbor.first = 0;
    neighbor.second = -1;
    neighborIterator.push_back(neighbor);
    neighbor.first = 1;
    neighbor.second = 0;
    neighborIterator.push_back(neighbor);

    allPushedIndX = new uint16_t[N_SCAN * Horizon_SCAN];
    allPushedIndY = new uint16_t[N_SCAN * Horizon_SCAN];

    queueIndX = new uint16_t[N_SCAN * Horizon_SCAN];
    queueIndY = new uint16_t[N_SCAN * Horizon_SCAN];

    gpfmethod = std::make_shared<GroundPlaneFit>();

}

//计算Scan之间点云的角度
void DrivableSeg::ComputeScanAngle(const cv::Mat &range_mat) {
    // todo 可以考虑一下角度图在时间上的关系
    angle_image = cv::Mat(N_SCAN, Horizon_SCAN, CV_32F, cv::Scalar::all(0.5 * M_PI));//默认角度为90度

    x_mat.row(0) = range_mat.row(0) * row_angle_cos[0];
    y_mat.row(0) = range_mat.row(0) * row_angle_sin[0];
    for (int r = 1; r < N_SCAN; r++) {
        x_mat.row(r) = range_mat.row(r) * row_angle_cos[r];
        y_mat.row(r) = range_mat.row(r) * row_angle_sin[r];
        for (int c = 0; c < Horizon_SCAN; c++) {
        if(range_mat.at<float>(r,c) == FLT_MAX || range_mat.at<float>(r-1,c) == FLT_MAX){
            angle_image.at<float>(r, c) = 0.5 * M_PI;
            continue;
        }
            dx = fabs(x_mat.at<float>(r - 1, c) - x_mat.at<float>(r, c));
            dy = fabs(y_mat.at<float>(r - 1, c) - y_mat.at<float>(r, c));
            angle_image.at<float>(r, c) = atan2(dy, dx);

        }
    }
}

// todo 选择合适的滤波方法
/*    cv::blur(angle_image,angle_image,cv::Size(3,1));
cv::medianBlur(angle_image,angle_image,3);
ApplySavitskyGolaySmoothing(angle_image,5);
cv::GaussianBlur(angle_image,angle_image,cv::Size(3,3),1);*/

void DrivableSeg::ComputeDrivableArea(pcl::PointCloud<pcl::PointXYZ>::ConstPtr laserCloudIn,
                                      pcl::PointCloud<pcl::PointXYZI> &segmentedCloud,
                                      const int cluster_purpose,
                                      const int minclusternub,
                                      bool &is_stair_flag) {
//    std::cout << "laserCloudIn cloud with " << laserCloudIn->size() << " points...\n";

    MinClusterNub = minclusternub;
    pcl::PointCloud<pcl::PointXYZI> CopylaserCloud;
    clusterpurpose = cluster_purpose;

    pcl::PointXYZI thisPoint;
    CopylaserCloud.clear();
    int Cloudsize = laserCloudIn->size();
    for(size_t i = 0;i<Cloudsize;i++){
        thisPoint.x = laserCloudIn->points[i].x;
        thisPoint.y = laserCloudIn->points[i].y;
        thisPoint.z = laserCloudIn->points[i].z;
        CopylaserCloud.push_back(thisPoint);
    }

    // todo 在投影的时候考虑加上膨胀
    ComputeRangeMat(CopylaserCloud, rangeMat);
//    ComputeScanAngle(rangeMat);
//    std::cout << "fullPointCloud cloud with " << fullPointCloud->size() << " points...\n";

    extract_object();

    if(clusterpurpose == 0){
        //障碍物的中心进行合并，主要针对椅子这种会过分割的物体，采用欧式聚类
        merge_object();
        segmentedCloud.clear();
        int cloudsize = SegObjectCloud->size();
        for(size_t j =0;j<cloudsize;j++){
            segmentedCloud.push_back(SegObjectCloud->points[j]);
        }
    }

    //提取平面的信息
    if(clusterpurpose == 1){
        if(maxclusterednub >= MinClusterNub) {//平面上的点满足数量要求
            pcl::PointCloud<pcl::PointXYZI>::Ptr PlaneCloud;
            PlaneCloud.reset(new pcl::PointCloud<pcl::PointXYZI>);
            ComputePlaneParam(PlaneCloud,is_stair_flag);
            segmentedCloud.clear();
            for(size_t j =0;j<PlaneCloud->size();j++){
                segmentedCloud.push_back(PlaneCloud->points[j]);
            }

        }
    }
}

void DrivableSeg::ComputePlaneParam(pcl::PointCloud<pcl::PointXYZI>::Ptr &PlaneCloud,
                                    bool &is_stair_flag){
    using Eigen::MatrixXf;
    using Eigen::JacobiSVD;
    using Eigen::VectorXf;
    using namespace std;
    using namespace cv;

    cv::Mat PlaneMat = cv::Mat::zeros(N_SCAN,Horizon_SCAN,CV_8S);

    for(size_t row = 1;row<N_SCAN;row++)
        for(size_t col = 0;col<Horizon_SCAN;col++){
            if(labelMat.at<int16_t>(row,col) == maxclusteredlabel) {
                PlaneMat.at<int8_t>(row, col) = 1;
                PlaneCloud->push_back(fullPointCloud->points[row * Horizon_SCAN + col]);
                PlaneCloud->back().intensity = 888;
            }
        }

    gpfmethod->setInputCloud(PlaneCloud);
    gpfmethod->filterGround(PlaneCloud,is_stair_flag);

}

void DrivableSeg::merge_object(){

    pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>);

    tree->setInputCloud(ClusteredCloudCenter);

    std::vector<pcl::PointIndices> cluster_indices;
    pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
    ec.setClusterTolerance(toler); // 2cm
    ec.setMinClusterSize(min_size);
    ec.setMaxClusterSize(max_size);
    ec.setSearchMethod(tree);
    ec.setInputCloud(ClusteredCloudCenter);
    ec.extract(cluster_indices);

    int cluster_index = 0;
    pcl::PointCloud<pcl::PointXYZI> TempCloud;
    for (std::vector<pcl::PointIndices>::const_iterator it = cluster_indices.begin();
         it != cluster_indices.end(); ++it) {
        cluster_index = clusteredcloudlabel[it->indices[0]];
        for (const auto &idx: it->indices) {
            for (int k = clusteredcloudstart[idx]; k < clusteredcloudstart[idx] + clusteredcloudnum[idx]; k++) {
                TempCloud.push_back(SegObjectCloud->points[k]);
                TempCloud.back().intensity = cluster_index;
//                SegObjectCloud->points[k].intensity = cluster_index;
            }
        }
    }
//    std::cout<<"labelCount"<<labelCount<<"cluster_index"<<cluster_index<<endl;
    SegObjectCloud->clear();
    pcl::copyPointCloud(TempCloud,*SegObjectCloud);


}

void DrivableSeg::ComputeRangeMat(const pcl::PointCloud <PointType> InputCloud, cv::Mat &range_mat) {
    float verticalAngle, horizonAngle, range;
    size_t rowIdn, columnIdn, index, cloudSize;
    PointType thisPoint;

    range_mat = cv::Mat(N_SCAN, Horizon_SCAN, CV_32F, cv::Scalar::all(FLT_MAX));
    cloudSize = InputCloud.points.size();
    std::fill(fullPointCloud->points.begin(), fullPointCloud->points.end(), nanPoint);

    int PointCnt = 0;
    for (size_t i = 0; i < cloudSize; ++i) {

        thisPoint.x = InputCloud.points[i].x;
        thisPoint.y = InputCloud.points[i].y;
        thisPoint.z = InputCloud.points[i].z;

        verticalAngle = atan2(thisPoint.z, sqrt(thisPoint.x * thisPoint.x + thisPoint.y * thisPoint.y)) * 180 / M_PI;
        rowIdn = (N_SCAN - 1) - (verticalAngle + vertical_bottom) / ang_res_y;

        if (rowIdn < 0 || rowIdn >= N_SCAN)
            continue;

        horizonAngle = atan2(thisPoint.y, thisPoint.x) * 180 / M_PI;

        columnIdn = (horizonAngle - horizon_bottom) / ang_res_x;

        if (columnIdn < 0 || columnIdn >= Horizon_SCAN)
            continue;

        //已经被投影过
        if (fullPointCloud->points[rowIdn * Horizon_SCAN + columnIdn].intensity != -1)
            continue;

        range = sqrt(thisPoint.x * thisPoint.x + thisPoint.y * thisPoint.y + thisPoint.z * thisPoint.z);
        if (range < sensorMinimumRange)
            continue;

        range_mat.at<float>(rowIdn, columnIdn) = range;

        index = rowIdn * Horizon_SCAN + columnIdn;
        thisPoint.intensity = range;
        fullPointCloud->points[index] = thisPoint;
        PointCnt++;

    }
    groundMat = cv::Mat::zeros(N_SCAN, Horizon_SCAN, CV_16S);
    for (int r = 0; r < N_SCAN; r++)
        for (int c = 0; c < Horizon_SCAN; c++) {
            if (fullPointCloud->points[r * Horizon_SCAN + c].intensity == -1)
                groundMat.at<int16_t>(r, c) = 1;//无效点标记为1
        }
}
void DrivableSeg::extract_object() {

    SegObjectCloud->clear();
    labelMat = groundMat.clone();

    clusteredcloudnum.clear();
    clusteredcloudstart.clear();
    clusteredcloudlabel.clear();
    ClusteredCloudCenter->clear();
    maxclusterednub = 0;
    maxclusteredlabel = 999;
    labelCount = 2;//
    for (size_t i = 0; i < N_SCAN; ++i)
        for (size_t j = 0; j < Horizon_SCAN; ++j)
            if (labelMat.at<int16_t>(i, j) == 0 && fullPointCloud->points[i*Horizon_SCAN+j].intensity != -1)//地面和没有深度的点的label是-1，其余是物体，对其进行标记
                labelComponents(i, j);

}


void DrivableSeg::labelComponents(int row, int col) {
    // use std::queue std::vector std::deque will slow the program down greatly
    float d1, d2, alpha, angle;
    int fromIndX, fromIndY, thisIndX, thisIndY;
    bool lineCountFlag[N_SCAN] = {false};

    static pcl::PointXYZ pointcloudmean;

    queueIndX[0] = row;//这个应该是待标记的队列
    queueIndY[0] = col;
    int queueSize = 1;
    int queueStartInd = 0;
    int queueEndInd = 1;

    allPushedIndX[0] = row;//这个应该是已经被标记了的队列
    allPushedIndY[0] = col;
    int allPushedIndSize = 1;
    //采用的BFS算法
    while (queueSize > 0) {
        // Pop point
        fromIndX = queueIndX[queueStartInd];
        fromIndY = queueIndY[queueStartInd];
        --queueSize;
        ++queueStartInd;
        // Mark popped point
        labelMat.at<int16_t>(fromIndX, fromIndY) = labelCount;

        // Loop through all the neighboring grids of popped grid
        for (auto iter = neighborIterator.begin(); iter != neighborIterator.end(); ++iter) {
            // new index
            thisIndX = fromIndX + (*iter).first;
            thisIndY = fromIndY + (*iter).second;

            // index should be within the boundary
            if (thisIndX < 0 || thisIndX >= N_SCAN)//X是行
                continue;
            // at range image margin (left or right side)
            if (thisIndY < 0 || thisIndY >= Horizon_SCAN)//Y是列
                continue;
            // prevent infinite loop (caused by put already examined point back)
            if (labelMat.at<int16_t>(thisIndX, thisIndY) != 0)//已经被标记过
                continue;

            d1 = std::max(rangeMat.at<float>(fromIndX, fromIndY),
                          rangeMat.at<float>(thisIndX, thisIndY));
            d2 = std::min(rangeMat.at<float>(fromIndX, fromIndY),
                          rangeMat.at<float>(thisIndX, thisIndY));

            if ((*iter).first == 0)//如果是左右相邻的点
                alpha = segmentAlphaY;
            else//如果是上下相邻的点
                alpha = segmentAlphaX;
            //这里角度的计算看不太明白
            angle = atan2(d2 * sin(alpha), (d1 - d2 * cos(alpha)));

            if (angle > segmentTheta){
                queueIndX[queueEndInd] = thisIndX;
                queueIndY[queueEndInd] = thisIndY;
                ++queueSize;
                ++queueEndInd;

                labelMat.at<int16_t>(thisIndX, thisIndY) = labelCount;
                lineCountFlag[thisIndX] = true;

                allPushedIndX[allPushedIndSize] = thisIndX;
                allPushedIndY[allPushedIndSize] = thisIndY;
                ++allPushedIndSize;
            }
        }
    }

    // check if this segment is valid
    bool feasibleSegment = false;
    if (allPushedIndSize >= MinClusterNub)
        feasibleSegment = true;
    else if (allPushedIndSize >= segmentValidPointNum && clusterpurpose == 0){//0表示对障碍物聚类
        int lineCount = 0;
        for (size_t i = 0; i < N_SCAN; ++i)
            if (lineCountFlag[i] == true)
                ++lineCount;
    if (lineCount >= segmentValidLineNum)
            feasibleSegment = true;
    }
    // segment is valid, mark these points
    if (feasibleSegment == true) {
        pointcloudmean.x = pointcloudmean.y = pointcloudmean.z = 0;
        clusteredcloudnum.push_back(allPushedIndSize);
        clusteredcloudstart.push_back(SegObjectCloud->size());
        clusteredcloudlabel.push_back(labelCount);
        for (size_t i = 0; i < allPushedIndSize; ++i) {
            SegObjectCloud->push_back(fullPointCloud->points[allPushedIndY[i] + allPushedIndX[i] * Horizon_SCAN]);
            SegObjectCloud->points.back().intensity = labelCount;//这个应该是设置分割后在rviz中显示的颜色的
            pointcloudmean.x += fullPointCloud->points[allPushedIndY[i] + allPushedIndX[i] * Horizon_SCAN].x;
            pointcloudmean.y += fullPointCloud->points[allPushedIndY[i] + allPushedIndX[i] * Horizon_SCAN].y;
            pointcloudmean.z += fullPointCloud->points[allPushedIndY[i] + allPushedIndX[i] * Horizon_SCAN].z;
        }
        pointcloudmean.x = (double)pointcloudmean.x/allPushedIndSize;
        pointcloudmean.y = (double)pointcloudmean.y/allPushedIndSize;
        pointcloudmean.z = (double)pointcloudmean.z/allPushedIndSize;
        ClusteredCloudCenter->push_back(pointcloudmean);
        if(allPushedIndSize > maxclusterednub){
            //用于进行提取平面的参数
            maxclusterednub = allPushedIndSize;
            maxclusteredlabel = labelCount;
        }
        ++labelCount;
    } else { // segment is invalid, mark these points
        for (size_t i = 0; i < allPushedIndSize; ++i)
            labelMat.at<int16_t>(allPushedIndX[i], allPushedIndY[i]) = -1;
    }
}


